from keras.datasets import mnist
from keras.utils import to_categorical
from keras.optimizers import Adam
import matplotlib.pyplot as plt
import numpy as np
import model
from model import Nets_2D
import os
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
from keras.callbacks import ReduceLROnPlateau,ModelCheckpoint,EarlyStopping

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))

sess = tf.Session()
sess.run(tf.global_variables_initializer())

# 仅第一次要进行下载，保存在对应python3/keras文件夹下
(x_train, y_train), (x_valid, y_valid) = mnist.load_data()

x_train = x_train[:,:,:,np.newaxis]
x_valid = x_valid[:,:,:,np.newaxis]
y_train = to_categorical(y_train) # one_hot编码
y_valid = to_categorical(y_valid)

print('x_train.shape: {}'.format(x_train.shape))
print('x_valid.shape: {}'.format(x_valid.shape))
print('y_train.shape: {}'.format(y_train.shape))
print('y_valid.shape: {}'.format(y_valid.shape))


# 参数设定
n_classes = 10
width     = 28
height    = 28
batch_size= 32
epochs    = 100
key       = 'mynet_2'

Nets = Nets_2D(width, height)
method = {
    'mynet_1' : Nets.mynet_1(),
    'mynet_2' : Nets.mynet_2()
}

m = method[key]

m.compile(loss='categorical_crossentropy',
          optimizer=Adam(lr=1.0e-3),
          metrics=['accuracy'])

callbacks = [
    ReduceLROnPlateau(monitor='accuracy', factor=0.5, patience=3, mode='max',
                       cooldown=1, min_lr=1e-10),
    ModelCheckpoint(key + 'checkpoints/_{}epochs.h5'.format(epochs),
                    monitor='accuracy',verbose=1,
                    save_best_only=True,mode='max'),
    EarlyStopping(monitor='accuracy',patience=1,verbose=1,mode='max')
]


hist = m.fit(x_train,y_train,validation_data=(x_valid,y_valid),
            batch_size=batch_size,epochs=epochs,verbose=1,
            callbacks=callbacks)